{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Problem: Implement Custom Loss Function (Huber Loss)\n",
    "\n",
    "### Problem Statement\n",
    "You are tasked with implementing the **Huber Loss** as a custom loss function in PyTorch. The Huber loss is a robust loss function used in regression tasks, less sensitive to outliers than Mean Squared Error (MSE). It transitions between L2 loss (squared error) and L1 loss (absolute error) based on a threshold parameter $ \\delta $.\n",
    "\n",
    "The Huber loss is mathematically defined as:\n",
    "$$\n",
    "L_{\\delta}(y, \\hat{y}) = \n",
    "\\begin{cases} \n",
    "\\frac{1}{2}(y - \\hat{y})^2 & \\text{for } |y - \\hat{y}| \\leq \\delta, \\\\\n",
    "\\delta \\cdot (|y - \\hat{y}| - \\frac{1}{2} \\delta) & \\text{for } |y - \\hat{y}| > \\delta,\n",
    "\\end{cases}\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $y$ is the true value,\n",
    "- $\\hat{y}$ is the predicted value,\n",
    "- $\\delta$ is a threshold parameter that controls the transition between L1 and L2 loss.\n",
    "\n",
    "### Requirements\n",
    "1. **Custom Loss Function**:\n",
    "   - Implement a class `HuberLoss` inheriting from `torch.nn.Module`.\n",
    "   - Define the `forward` method to compute the Huber loss as per the formula.\n",
    "\n",
    "2. **Usage in a Regression Model**:\n",
    "   - Integrate the custom loss function into a regression training pipeline.\n",
    "   - Use it to compute and optimize the loss during model training.\n",
    "\n",
    "### Constraints\n",
    "- The implementation must handle both scalar and batch inputs for $ y $ (true values) and $ \\hat{y} $ (predicted values).\n",
    "\n",
    "\n",
    "Extra Details: https://en.wikipedia.org/wiki/Huber_loss\n",
    "\n",
    "<details>\n",
    "  <summary>💡 Hint</summary>\n",
    "  Some details: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch/notebook\n",
    "</details>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate synthetic data\n",
    "torch.manual_seed(42)\n",
    "X = torch.rand(100, 1) * 10  # 100 data points between 0 and 10\n",
    "y = 2 * X + 3 + torch.randn(100, 1)  # Linear relationship with noise\n",
    "\n",
    "\n",
    "class HuberLoss(nn.Module):\n",
    "    def __init__(self, delta=1.0):\n",
    "        super(HuberLoss, self).__init__()\n",
    "        self.delta = delta\n",
    "    \n",
    "    def forward(self, y_pred, y_true):\n",
    "        # Calculate the absolute error\n",
    "        error = torch.abs(y_pred - y_true)\n",
    "        \n",
    "        # Apply the Huber loss formula\n",
    "        loss = torch.where(error <= self.delta,\n",
    "                           0.5 * error**2,  # L2 loss for small errors\n",
    "                           self.delta * (error - 0.5 * self.delta))  # L1 loss for large errors\n",
    "        return loss.mean()  # Return the mean loss across all samples\n",
    "\n",
    "\n",
    "# Define the Linear Regression Model\n",
    "class LinearRegressionModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LinearRegressionModel, self).__init__()\n",
    "        self.linear = nn.Linear(1, 1)  # Single input and single output\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "# Initialize the model, loss function, and optimizer\n",
    "model = LinearRegressionModel()\n",
    "criterion = HuberLoss(delta=1.0)\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01)\n",
    "\n",
    "# Training loop\n",
    "epochs = 1000\n",
    "for epoch in range(epochs):\n",
    "    # Forward pass\n",
    "    predictions = model(X)\n",
    "    loss = criterion(predictions, y)\n",
    "\n",
    "    # Backward pass and optimization\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # Log progress every 100 epochs\n",
    "    if (epoch + 1) % 100 == 0:\n",
    "        print(f\"Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the learned parameters\n",
    "[w, b] = model.linear.parameters()\n",
    "print(f\"Learned weight: {w.item():.4f}, Learned bias: {b.item():.4f}\")\n",
    "\n",
    "# Testing on new data\n",
    "X_test = torch.tensor([[4.0], [7.0]])\n",
    "with torch.no_grad():\n",
    "    predictions = model(X_test)\n",
    "    print(f\"Predictions for {X_test.tolist()}: {predictions.tolist()}\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
